# python的flask网页版聊天程序web-ui，调用zhipuAI大模型进行人机对话。 
# 采用大模型的chat_models模式，用户输入问题，大模型给出问答，然后将AI和Human的Message压缩后作为上下文，进行下一轮对话生成。 
# 用户输入/quit对话结束时，将Message历史保存为sqlite数据库。

from flask import Flask, render_template, request, jsonify
import requests
import sqlite3
import json
import os


# 智谱AI的 API 地址和密钥
from dotenv import load_dotenv
load_dotenv()  # 加载环境变量
API_KEY = os.getenv("ZHIPUAI_API_KEY")  # 从环境变量中获取 API 密钥
API_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
model_name = "glm-4-flash"  # 使用的模型



dir_path = os.path.dirname(os.path.realpath(__file__))
db_path = os.path.join(dir_path, 'chat_history.db')

app = Flask(__name__)

# SQLite 数据库初始化
def init_db():

    # conn = sqlite3.connect('chat_history.db')
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS chat_history
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  user_input TEXT,
                  ai_response TEXT)''')
    conn.commit()
    conn.close()

# 调用智谱AI的 chat_models 模式
def call_zhipuai(messages):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json",
    }
    data = {
        "model": model_name,  # 使用的模型
        "messages": messages,    # 对话上下文
        "stream": False           # 非流式返回
    }
    response = requests.post(API_URL, headers=headers, json=data)
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"Error: {response.status_code} - {response.text}"

# 保存对话历史到 SQLite 数据库
def save_chat_history(user_input, ai_response):
    # conn = sqlite3.connect('chat_history.db')  # 全局的数据库地址
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("INSERT INTO chat_history (user_input, ai_response) VALUES (?, ?)",
              (user_input, ai_response))
    conn.commit()
    conn.close()

# 主页路由
@app.route('/')
def index():
    return render_template('index_app01.html')

# 聊天路由
@app.route('/chat', methods=['POST'])
def chat():
    user_input = request.json.get('message')
    if user_input == "/quit":
        return jsonify({"response": "对话已结束，历史记录已保存。"})

    # 获取当前对话上下文
    session_messages = request.json.get('messages', [])
    session_messages.append({"role": "user", "content": user_input})

    # 调用智谱AI
    ai_response = call_zhipuai(session_messages)
    session_messages.append({"role": "assistant", "content": ai_response})

    # 保存对话历史
    save_chat_history(user_input, ai_response)

    # 返回 AI 的响应和更新后的上下文
    return jsonify({"response": ai_response, "messages": session_messages})

if __name__ == '__main__':
    init_db()  # 初始化数据库
    app.run(debug=True)
